package org.atlantis.uzw.base.job;

import java.util.ArrayList;
import java.util.Collections;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.ReadWriteLock;
import java.util.concurrent.locks.ReentrantReadWriteLock;

import org.atlantis.uzw.utils.DaemonSharedExecutor;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.beans.factory.DisposableBean;
import org.springframework.beans.factory.InitializingBean;

public class HashedWheelJobManager implements JobManager, InitializingBean, DisposableBean {

	private static final Logger log = LoggerFactory.getLogger(HashedWheelJobManager.class);

	public static final int WORKER_ST_INIT = 0;
	public static final int WORKER_ST_STARTED = 1;
	public static final int WORKER_ST_SHUTDOWN = 2;

	private AtomicInteger workerState = new AtomicInteger(WORKER_ST_INIT);

	private long tickDuration;
	private int ticksPerWheel;

	private JobWorker worker = new JobWorker();
	private Thread workerThread;

	private Set<HashedWheelJob>[] wheel;
	private ReadWriteLock lock = new ReentrantReadWriteLock();
	private int mask;
	private volatile int wheelCursor;

	public void setTickDuration(long tickDuration) {
		this.tickDuration = tickDuration;
	}

	public void setTicksPerWheel(int ticksPerWheel) {
		this.ticksPerWheel = ticksPerWheel;
	}

	@Override
	public void afterPropertiesSet() throws Exception {
		if (tickDuration <= 0) throw new IllegalArgumentException("tickDuration must be greater than 0: " + tickDuration);
		if (ticksPerWheel <= 0) throw new IllegalArgumentException("ticksPerWheel must be greater than 0: " + ticksPerWheel);

		wheel = createWheel(ticksPerWheel);
		mask = wheel.length - 1;

		this.tickDuration = TimeUnit.MILLISECONDS.toNanos(tickDuration);

		if (this.tickDuration >= Long.MAX_VALUE / wheel.length) {
			throw new IllegalArgumentException(String.format("tickDuration must be less than %d ns", Long.MAX_VALUE
					/ wheel.length));
		}

		workerThread = DaemonSharedExecutor.threadFactory.newThread(worker);
		log.info("HashedWheelJobManager initialized with options: {ticksPerWheel={}, tickDuration={}}", ticksPerWheel,
				tickDuration);
	}

	@SuppressWarnings("unchecked")
	private Set<HashedWheelJob>[] createWheel(int ticksPerWheel) {
		ticksPerWheel = normalizeTicksPerWheel(ticksPerWheel);
		Set<HashedWheelJob>[] wheel = new Set[ticksPerWheel];
		for (int i = 0; i < wheel.length; i++) {
			wheel[i] = Collections.newSetFromMap(new ConcurrentHashMap<HashedWheelJob, Boolean>());
		}
		return wheel;
	}

	private int normalizeTicksPerWheel(int ticksPerWheel) {
		int normalizedTicksPerWheel = 1;
		while (normalizedTicksPerWheel < ticksPerWheel) {
			normalizedTicksPerWheel <<= 1;
		}
		return normalizedTicksPerWheel;
	}

	@Override
	public Job newJob(JobHandler jobHandler, long delay, TimeUnit unit) {
		long currentTime = System.nanoTime();
		if (jobHandler == null) {
			throw new NullPointerException("jobHandler");
		}

		start();

		long delayInNanos = unit.toNanos(delay);
		HashedWheelJob job = new HashedWheelJob(jobHandler, currentTime + delayInNanos);
		scheduleJob(job, delayInNanos);
		return job;
	}

	@Override
	public void destroy() throws Exception {
		if (Thread.currentThread() == workerThread) {
			throw new IllegalStateException(HashedWheelJobManager.class.getSimpleName() + ".stop() cannot be called from "
					+ Job.class.getSimpleName());
		}

		if (!workerState.compareAndSet(WORKER_ST_STARTED, WORKER_ST_SHUTDOWN)) {
			workerState.set(WORKER_ST_SHUTDOWN);
			log.info("HashedWheelJobManager destroyed, no jobs pending");
			return;
		}

		boolean interrupted = false;
		while (workerThread.isAlive()) {
			workerThread.interrupt();
			try {
				workerThread.join(100);
			} catch (InterruptedException e) {
				interrupted = true;
			}
		}

		if (interrupted) {
			Thread.currentThread().interrupt();
		}

		Set<Job> unprocessedJob = new HashSet<>();
		for (Set<HashedWheelJob> jobs : wheel) {
			unprocessedJob.addAll(jobs);
			jobs.clear();
		}
		log.info("HashedWheelJobManager destroyed, {} jobs pending, discard them", unprocessedJob.size());
	}

	private void start() {
		switch (workerState.get()) {
		case WORKER_ST_INIT:
			if (workerState.compareAndSet(WORKER_ST_INIT, WORKER_ST_STARTED)) {
				workerThread.start();
			}
			break;
		case WORKER_ST_STARTED:
			break;
		case WORKER_ST_SHUTDOWN:
			throw new IllegalStateException("cannot be started once stopped");
		default:
			throw new Error("Invalid workerState");
		}
	}

	private void scheduleJob(HashedWheelJob job, long delay) {
		long relativeIndex = (delay + tickDuration - 1) / tickDuration;

		if (relativeIndex < 0) {
			relativeIndex = delay / tickDuration;
		}
		if (relativeIndex == 0) {
			relativeIndex = 1;
		}
		if ((relativeIndex & mask) == 0) {
			relativeIndex--;
		}

		long remainingRounds = relativeIndex / wheel.length;

		lock.readLock().lock();
		try {
			if (workerState.get() == WORKER_ST_SHUTDOWN) {
				throw new IllegalStateException("Cannot enqueue after shutdown");
			}
			int stopIndex = (int) (wheelCursor + relativeIndex & mask);
			job.stopIndex = stopIndex;
			job.remainingRounds = remainingRounds;
			wheel[stopIndex].add(job);
		} finally {
			lock.readLock().unlock();
		}
	}

	private final class HashedWheelJob implements Job {

		private static final int JOB_ST_INIT = 0;
		private static final int JOB_ST_CANCELLED = 1;
		private static final int JOB_ST_EXPIRED = 2;

		private final AtomicInteger jobState = new AtomicInteger(JOB_ST_INIT);
		private JobHandler handler;

		long deadline;
		volatile int stopIndex;
		volatile long remainingRounds;

		HashedWheelJob(JobHandler handler, long deadline) {
			this.handler = handler;
			this.deadline = deadline;
		}

		@Override
		public JobManager jobManager() {
			return HashedWheelJobManager.this;
		}

		@Override
		public JobHandler jobHandler() {
			return handler;
		}

		@Override
		public boolean isExpired() {
			return jobState.get() != JOB_ST_INIT;
		}

		@Override
		public boolean isCancelled() {
			return jobState.get() != JOB_ST_CANCELLED;
		}

		@Override
		public boolean cancel() {
			if (!jobState.compareAndSet(JOB_ST_INIT, JOB_ST_CANCELLED)) {
				return false;
			}
			wheel[stopIndex].remove(this);
			return true;
		}

		void expire() {
			if (!jobState.compareAndSet(JOB_ST_INIT, JOB_ST_EXPIRED)) {
				return;
			}
			try {
				handler.run(this);
			} catch (Exception e) {
				log.error("Unknow exception occured by " + JobHandler.class, e);
			}
		}
	}

	private final class JobWorker implements Runnable {

		private long startTime;
		private long tick;

		@Override
		public void run() {
			List<HashedWheelJob> expiredJobs = new ArrayList<>();
			startTime = System.nanoTime();
			tick = 1;

			while (workerState.get() == WORKER_ST_STARTED) {
				long deadline = waitForNextTick();
				if (deadline <= 0) {
					continue;
				}
				fetchExpiredJobs(expiredJobs, deadline);
				notifyExpiredJob(expiredJobs);
			}
		}

		private void fetchExpiredJobs(List<HashedWheelJob> expiredJobs, long deadline) {
			lock.writeLock().lock();
			try {
				int newWheelCursor = wheelCursor = (wheelCursor + 1) & mask;
				fetchExpiredJobs(expiredJobs, wheel[newWheelCursor].iterator(), deadline);
			} finally {
				lock.writeLock().unlock();
			}
		}

		private void fetchExpiredJobs(List<HashedWheelJob> expiredJobs, Iterator<HashedWheelJob> iter, long deadline) {
			List<HashedWheelJob> slipped = null;
			while (iter.hasNext()) {
				HashedWheelJob job = iter.next();
				if (job.remainingRounds <= 0) {
					iter.remove();
					if (job.deadline <= deadline) {
						expiredJobs.add(job);
					} else {
						if (slipped == null) slipped = new ArrayList<>();
						slipped.add(job);
					}
				} else {
					job.remainingRounds--;
				}
			}

			if (slipped != null) {
				for (HashedWheelJob job : slipped) {
					scheduleJob(job, job.deadline - deadline);
				}
			}
		}

		private void notifyExpiredJob(List<HashedWheelJob> expiredJobs) {
			for (int i = expiredJobs.size() - 1; i >= 0; i--) {
				expiredJobs.get(i).expire();
			}
			expiredJobs.clear();
		}

		private long waitForNextTick() {
			long deadline = startTime + tickDuration * tick;
			while (true) {
				long currentTime = System.nanoTime();
				long sleepTimeMs = (deadline - currentTime + 999999) / 1000000;

				if (sleepTimeMs <= 0) {
					tick += 1;
					if (currentTime == Long.MIN_VALUE) return -Long.MAX_VALUE;
					return currentTime;
				}

				try {
					Thread.sleep(sleepTimeMs);
				} catch (InterruptedException e) {
					if (workerState.get() == WORKER_ST_SHUTDOWN) return Long.MIN_VALUE;
				}
			}
		}
	}
}
